from collections import defaultdict
from typing import Literal
from conn import mysql_stage_engine, clip_col, clip_annotated_record, mysql_engine
from sqlalchemy import Column, Integer, String, and_, delete, text
from sqlalchemy.orm import declarative_base, Session
from loguru import logger
from datetime import datetime, timedelta
from tqdm import tqdm
import pytz


SH_TZ = pytz.timezone("Asia/Shanghai")
Base = declarative_base()

supplier_dict = {
    0: "智目",
    1: "恺望",
    2: "百度",
    3: "ZDriveAI",
    4: "文远知行",
    5: "腾讯",
    6: "柏川",
    7: "曼孚",
    8: "百度4D",
    9: "整数",
    11: "澳鹏",
    10: "数据堂",
    12: "标贝",
    13: "ZdriveLabel",
    14: "曼孚-2",
}


CntByClientType = dict[Literal["sensetime", "imotion", "zdrive"], int]
PeriodType = Literal["day", "week", "month", "quarter", "year"]
ProjectType = Literal["driving", "parking"]


CLIENT_LIST = [
    dict(name="sensetime", zh_name="商汤"),
    dict(name="imotion", zh_name="知行"),
    dict(name="zdrive", zh_name="大卓"),
]
POOL_SIZE = 10

QUERY_BATCH_SIZE = 1000


def has_parking_tag(tags: dict[str, list[str]]):
    route_tags = tags.get("route") or []
    return "indoor_parking" in route_tags or "open_parking" in route_tags


def has_driving_tag(tags: dict[str, list[str]]):
    # TODO 再确认一下这个逻辑
    route_tags = tags.get("route") or []
    return "highway" in route_tags or (
        "city_road" in route_tags
        and "open_parking" not in route_tags
        and "indoor_parking" not in route_tags
    )


# 行车和泊车对应的标注类型
# 23 的话需要 tag 满足 route-city_road
# TODO 需求量先写死
driving_annotate_types_dict = {
    2: {
        "name": "行车 车道线",
        "demand_clip_cnt": 75_0000,
        "valid_tag": lambda tags: True,
    },
    5: {
        "name": "行车 红绿灯",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    15: {
        "name": "行车 PVB-城区",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    16: {
        "name": "行车 PVB-高速",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    22: {
        "name": "行车 OCC",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    23: {"name": "行车 GOP", "demand_clip_cnt": 50_0000, "valid_tag": has_driving_tag},
    24: {
        "name": "行车 相机遮挡",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_driving_tag,
    },
    25: {
        "name": "行车 环境识别",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_driving_tag,
    },
    27: {
        "name": "行车 交通标志",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    28: {
        "name": "行车 光斑",
        "demand_clip_cnt": 30_0000,
        "valid_tag": lambda tags: True,
    },
    32: {
        "name": "行车 3D-OCC",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
}
driving_annotate_types = list(driving_annotate_types_dict.keys())


# 23 的话需要 tag 满足 route-indoor_parking route-open_parking
parking_annotate_types_dict = {
    17: {
        "name": "泊车 PVB",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    19: {
        "name": "泊车 车位",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    20: {
        "name": "泊车 路面标识",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
    21: {
        "name": "泊车 OCC",
        "demand_clip_cnt": 100_0000,
        "valid_tag": lambda tags: True,
    },
    23: {"name": "泊车 GOP", "demand_clip_cnt": 50_0000, "valid_tag": has_parking_tag},
    24: {
        "name": "泊车 相机遮挡",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_parking_tag,
    },
    25: {
        "name": "泊车 环境识别",
        "demand_clip_cnt": 50_0000,
        "valid_tag": has_parking_tag,
    },
    31: {
        "name": "泊车 车道线",
        "demand_clip_cnt": 50_0000,
        "valid_tag": lambda tags: True,
    },
}

parking_annotate_types = list(parking_annotate_types_dict.keys())

annotate_types_by_project = {
    "parking": parking_annotate_types_dict,
    "driving": driving_annotate_types_dict,
}

SH_TZ = pytz.timezone("Asia/Shanghai")


def orm_to_dict(obj):
    return {c.name: getattr(obj, c.name) for c in obj.__table__.columns}


class CurateDashboardORM(Base):
    __tablename__ = "curate_dashboard"

    id = Column(Integer, primary_key=True, autoincrement=True)
    project = Column(String(255), nullable=False, comment="项目名称")
    period = Column(
        String(10),
        nullable=False,
        comment="统计周期",
    )
    annotate_type = Column(Integer, nullable=False, comment="标注类型")
    annotate_name = Column(String(255), nullable=False, comment="标注类型名称")
    supplier = Column(String(255), nullable=False, comment="供应商")
    start_date = Column(String(10), nullable=False, comment="开始日期")
    end_date = Column(String(10), nullable=False, comment="结束日期")

    curated_clip_cnt = Column(Integer, default=0, comment="送标 clip 数")
    curated_frame_cnt = Column(Integer, default=0, comment="送标 frame 数")
    labeled_clip_cnt = Column(Integer, default=0, comment="回标 clip 数")
    labeled_frame_cnt = Column(Integer, default=0, comment="回标 frame 数")
    finished_clip_cnt = Column(Integer, default=0, comment="确认 clip 数")
    finished_frame_cnt = Column(Integer, default=0, comment="确认 frame 数")


# 回标统计
def handle_labeled_fields(
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    logger.info(
        f"handle_labeled_fields for project: {project}, annotate_type: {annotate_type}, date_range: [{start_date}, {end_date})"
    )
    start_ts = int(datetime.strptime(start_date, "%Y-%m-%d").timestamp() * 1000)
    end_ts = int(datetime.strptime(end_date, "%Y-%m-%d").timestamp() * 1000)

    filter = dict(
        ctime={"$gte": start_ts, "$lt": end_ts},
        annotate_type=annotate_type,
    )
    rows = []  # [supplier, clip_cnt, frame_cnt]

    supplier_data_cnt_map = {}
    clip_records = clip_annotated_record.find(
        filter, {"tags": 1, "frame_count": 1, "supplier": 1}
    ).to_list()
    for clip_record in clip_records:
        if clip_record["supplier"] not in supplier_data_cnt_map:
            supplier_data_cnt_map[clip_record["supplier"]] = {
                "frame_cnt": 0,
                "clip_cnt": 0,
            }
        if annotate_types_by_project[project][annotate_type]["valid_tag"](
            clip_record["tags"]
        ):
            supplier_data_cnt_map[clip_record["supplier"]]["frame_cnt"] += clip_record[
                "frame_count"
            ]
            supplier_data_cnt_map[clip_record["supplier"]]["clip_cnt"] += 1
    for supplier, data_cnt in supplier_data_cnt_map.items():
        rows.append((supplier, data_cnt["clip_cnt"], data_cnt["frame_cnt"]))
    return rows


# 确认统计
def handle_finished_fields(
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    logger.info(
        f"handle_finished_field for project: {project}, annotate_type: {annotate_type}, date_range: [{start_date}, {end_date})"
    )
    start_ts = int(datetime.strptime(start_date, "%Y-%m-%d").timestamp() * 1000)
    end_ts = int(datetime.strptime(end_date, "%Y-%m-%d").timestamp() * 1000)
    filter = dict(
        confirmed_time={"$gte": start_ts, "$lt": end_ts},
        state=2,  # 确认
        annotate_type=annotate_type,
    )
    rows = []
    if annotate_type != 23:
        groups = clip_annotated_record.aggregate(
            [
                {"$match": filter},
                {
                    "$group": {
                        "_id": "$supplier",
                        "clip_cnt": {"$sum": 1},
                        "frame_cnt": {"$sum": "$frame_count"},
                    }
                },
            ]
        )
        for group in groups:
            rows.append((group["_id"], group["clip_cnt"], group["frame_cnt"]))
        return rows

    supplier_data_cnt_map = {}
    clip_records = clip_annotated_record.find(
        filter, {"tags": 1, "frame_count": 1, "supplier": 1}
    ).to_list()
    for clip_record in clip_records:
        if clip_record["supplier"] not in supplier_data_cnt_map:
            supplier_data_cnt_map[clip_record["supplier"]] = {
                "frame_cnt": 0,
                "clip_cnt": 0,
            }
        if annotate_types_by_project[project][annotate_type]["valid_tag"](
            clip_record["tags"]
        ):
            supplier_data_cnt_map[clip_record["supplier"]]["frame_cnt"] += clip_record[
                "frame_count"
            ]
            supplier_data_cnt_map[clip_record["supplier"]]["clip_cnt"] += 1
    for supplier, data_cnt in supplier_data_cnt_map.items():
        rows.append((supplier, data_cnt["clip_cnt"], data_cnt["frame_cnt"]))
    return rows


# 送标统计
def handle_curated_fields(
    project: ProjectType,
    annotate_type: int,
    start_date: str,
    end_date: str,
):
    logger.info(
        f"handle_curated_field for project: {project}, annotate_type: {annotate_type}, date_range: [{start_date}, {end_date})"
    )
    start_ts = int(datetime.strptime(start_date, "%Y-%m-%d").timestamp() * 1000)
    end_ts = int(datetime.strptime(end_date, "%Y-%m-%d").timestamp() * 1000)

    rows = []  # supplier, clip_cnt, frame_cnt

    supplier_data_cnt_map = {}
    with mysql_engine.connect() as conn:
        sql = text(
            """
            select id, supplier from send_label_task_tab 
            where ctime >= :start_ts and ctime < :end_ts 
            and annotate_type = :annotate_type
            """
        ).bindparams(start_ts=start_ts, end_ts=end_ts, annotate_type=annotate_type)

        all_task = conn.execute(sql).fetchall()
        for task_id, supplier in tqdm(all_task):
            if supplier not in supplier_data_cnt_map:
                supplier_data_cnt_map[supplier] = {"frame_cnt": 0, "clip_cnt": 0}
            sql = text(
                """
                select clip_id, count(distinct frame_id)
                from send_label_record_tab 
                where send_label_task_id = :task_id
                group by clip_id
                """
            ).bindparams(task_id=task_id)
            # NOTE clip 数量可能很多，只能分批加载数据
            all_cnt = conn.execute(sql).fetchall()
            for i in range(0, len(all_cnt), 1000):
                clip_ids = [clip_id for clip_id, _ in all_cnt[i : i + 1000]]
                clip_tags = clip_col.find(
                    {"_id": {"$in": clip_ids}},
                    {"tags": 1},
                )
                clip_tags_map = {item["_id"]: item for item in clip_tags}

                for clip_id, frame_cnt in all_cnt[i : i + 1000]:
                    if annotate_types_by_project[project][annotate_type]["valid_tag"](
                        clip_tags_map[clip_id]["tags"]
                    ):
                        supplier_data_cnt_map[supplier]["clip_cnt"] += 1
                        supplier_data_cnt_map[supplier]["frame_cnt"] += frame_cnt
    for supplier, data_cnt in supplier_data_cnt_map.items():
        rows.append((supplier, data_cnt["clip_cnt"], data_cnt["frame_cnt"]))
    return rows


def build_dashboard_by_project(
    project: ProjectType, period: str, start_date: str, end_date: str
):
    dashboard_by_supplier_anno_type = defaultdict(
        lambda: CurateDashboardORM(
            project=project,
            period=period,
            start_date=start_date,
            end_date=end_date,
            annotate_type=0,  # fill later
            annotate_name="",  # fill later
            supplier="",  # fill later
            curated_clip_cnt=0,
            curated_frame_cnt=0,
            labeled_clip_cnt=0,
            labeled_frame_cnt=0,
            finished_clip_cnt=0,
            finished_frame_cnt=0,
        )
    )
    annotate_types = annotate_types_by_project[project]
    for annotate_type in annotate_types:
        supplier_clip_frame_cnt = handle_curated_fields(
            project, annotate_type, start_date, end_date
        )
        for supplier, clip_cnt, frame_cnt in supplier_clip_frame_cnt:
            key = (annotate_type, supplier)
            dashboard_by_supplier_anno_type[key].curated_clip_cnt += clip_cnt
            dashboard_by_supplier_anno_type[key].curated_frame_cnt += frame_cnt

        supplier_clip_frame_cnt = handle_labeled_fields(
            project, annotate_type, start_date, end_date
        )
        for supplier, clip_cnt, frame_cnt in supplier_clip_frame_cnt:
            key = (annotate_type, supplier)
            dashboard_by_supplier_anno_type[key].labeled_clip_cnt += clip_cnt
            dashboard_by_supplier_anno_type[key].labeled_frame_cnt += frame_cnt

        supplier_clip_frame_cnt = handle_finished_fields(
            project, annotate_type, start_date, end_date
        )
        for supplier, clip_cnt, frame_cnt in supplier_clip_frame_cnt:
            key = (annotate_type, supplier)
            dashboard_by_supplier_anno_type[key].finished_clip_cnt += clip_cnt
            dashboard_by_supplier_anno_type[key].finished_frame_cnt += frame_cnt

    dashboard_list = []
    for (annotate_type, supplier), instance in dashboard_by_supplier_anno_type.items():
        instance.annotate_type = annotate_type
        instance.annotate_name = annotate_types_by_project[project][annotate_type][
            "name"
        ]
        instance.supplier = supplier_dict[int(supplier)]
        dashboard_list.append(instance)
    return dashboard_list


def build_dashboard(period: str, start_date: str, end_date: str):
    driving_data_list = build_dashboard_by_project(
        "driving", period, start_date, end_date
    )
    parking_data_list = build_dashboard_by_project(
        "parking", period, start_date, end_date
    )
    return [*driving_data_list, *parking_data_list]


def insert_dashboard(dashboard_list: list[CurateDashboardORM]):
    with Session(mysql_stage_engine) as session:
        session.add_all(dashboard_list)
        session.commit()


def delete_dashboard(period: str, start_date: str):
    with Session(mysql_stage_engine) as session:
        del_stmt = delete(CurateDashboardORM).where(
            and_(
                CurateDashboardORM.period == period,
                CurateDashboardORM.start_date == start_date,
            )
        )
        ret = session.execute(del_stmt)
        session.commit()
    return ret.rowcount


def build_dashboard_by_date(date: str | None = None):
    if date is None:
        # 自动运行，计算的是前一天的数据
        end_date = datetime.now().astimezone(SH_TZ).date()
        start_date = end_date - timedelta(days=1)
    else:
        # 用户指定日期，则计算当天数据
        start_date = datetime.strptime(date, "%Y-%m-%d").astimezone(SH_TZ).date()
        end_date = start_date + timedelta(days=1)

    start_date_str, end_date_str = (
        start_date.strftime("%Y-%m-%d"),
        end_date.strftime("%Y-%m-%d"),
    )

    # 每天执行一次，计算前一天的数据
    data_list = build_dashboard("day", start_date_str, end_date_str)
    delete_row_cnt = delete_dashboard("day", start_date_str)
    logger.info(
        f"delete day dashboard for date range [{start_date_str}, {end_date_str}): {delete_row_cnt}"
    )
    insert_dashboard(data_list)
    logger.info(
        f"insert day dashboard for date range [{start_date_str}, {end_date_str}): {len(data_list)}"
    )

    # 每天执行一次，计算累计数据
    if end_date == datetime.now().astimezone(SH_TZ).date():
        start_date = (
            datetime.strptime("2023-01-01", "%Y-%m-%d").astimezone(SH_TZ).date()
        )
        start_date_str = start_date.strftime("%Y-%m-%d")
        logger.info(
            f"build total dashboard for date range [{start_date_str}, {end_date_str})"
        )

        data_list = build_dashboard("total", start_date_str, end_date_str)
        delete_row_cnt = delete_dashboard("total", start_date_str)
        logger.info(
            f"delete total dashboard for date range [{start_date_str}, {end_date_str}): {delete_row_cnt}"
        )
        insert_dashboard(data_list)
        logger.info(
            f"insert total dashboard for date range [{start_date_str}, {end_date_str}): {len(data_list)}"
        )

    if end_date.weekday() == 0:
        # 每周一执行一次
        start_date = end_date - timedelta(days=7)
        start_date_str = start_date.strftime("%Y-%m-%d")
        data_list = build_dashboard("week", start_date_str, end_date_str)
        logger.info(
            f"delete week dashboard for date range[{start_date_str}, {end_date_str}): {delete_row_cnt}"
        )
        delete_row_cnt = delete_dashboard("week", start_date_str)
        insert_dashboard(data_list)
        logger.info(
            f"insert week dashboard for date range[{start_date_str}, {end_date_str}): {len(data_list)}"
        )


if __name__ == "__main__":
    build_dashboard_by_date()
